#!/bin/env python3
import ithildin as ith
import numpy as np
import matplotlib.pyplot as plt
import sys
from typing import List
from scipy.interpolate import RectBivariateSpline
from pydiffmap.diffusion_map import DiffusionMap
from pydiffmap import visualization as diff_visualization
import matplotlib

# Courtemanche1998 model, 1 ruimtedimensies, 21 variabelen

data = ith.SimData.from_stem("myokit12/myokit_1")

# Vermijd randeffecten en initialiatie
# TODO: begintijd zou in principe in de log moeten staan
def snij_randen_weg(variabele,begintijd, eindtijd):
    def begin(index):
        if index == 0:
            return begintijd
        return variabele.shape[index]//8
    def eind(index):
        if index == 0:
            return eindtijd
        if variabele.shape[index] < 20:
            return variabele.shape[index]
        return (7 * variabele.shape[index])//8
    return variabele[begin(0):eind(0),begin(1):eind(1),begin(2):eind(2),begin(3):eind(3)]

def snij_randen_weg2(variables,begintijd,eindtijd):
    newdata = dict()
    for key in variables.keys():
        newdata[key] = snij_randen_weg(variables[key], begintijd,eindtijd)
    return newdata

def onderzoek(t0, t1, aantal_punten, eps, bandwidth_normalize=False, normalize=False, k=64,seed=1):
    vars = snij_randen_weg2(data.vars, t0, t1)
    # We bekijken de faseruimte, tijds- en ruimtecoordinaten zijn daarin niet van belang.
    for key in vars.keys():
        vars[key] = vars[key].ravel()
    # Verlaag aantal punten om geheugengebruik en tijdsduur te beperken.
    # replace=True strikt genomen incorrect maar verlaagt geheugengebruik
    # en wegens de hoge hoeveelheid punten weinig belangrijk
    np.random.seed(seed)
    keuzes = np.random.choice(vars['u'].shape[0], aantal_punten,replace=True)
    for key in data.vars.keys():
        vars[key] = vars[key][keuzes]
        if normalize==True and np.std(vars[key]) != 0: # K is constant
            vars[key] = vars[key] - np.mean(vars[key])
            vars[key] = vars[key] / np.std(vars[key])
    metric = 'euclidean'
    if normalize=='mahalanobis':
        metric = 'mahalanobis'
        del vars['K']
    # Nu de randen verwijderd zijn en we punten in een faseruimte hebben, kunnens we proberen
    # diffusion maps te gebruiken.
    dmap = DiffusionMap.from_sklearn(epsilon=eps,n_evecs=20,bandwidth_normalize=bandwidth_normalize,metric=metric,k=k)
    # Don't use K because that variable is constant
    dmap.fit(np.column_stack((vars['gateui'], vars['gatexs'], vars['gated'], vars['Ca'], vars['gatew'], vars['gateu'], vars['Na'], vars['gateoa'], vars['Caup'], vars['gatef'], vars['gateua'], vars['gateh'], vars['gatexr'], vars['gatefCa'], vars['u'], vars['gatem'], vars['gatev'], vars['gatej'], vars['gateoi'], vars['Carel'])))
    scatter_kwargs = { 'c' : dmap.dmap[:,3]}
    diff_visualization.embedding_plot(dmap,dim=3,scatter_kwargs=scatter_kwargs)
    return dmap

# Zoals ‘Embedding given by first three DCs, coloured by fourth (BOFC model, %s points)’,
# maar plot ook de punten die niet gebruikt worden om de diffusion map op te stellen.
# Hopelijk maakt dat de figuur wat duidelijker ...
def plot_veel_punten(t0, t1, aantal_punten, dmap):
    vars = snij_randen_weg2(data.vars, t0, t1)
    # copied from 'onderzoek'
    for key in vars.keys():
        vars[key] = vars[key].ravel()
    print("raveled")
    np.random.seed(1) # determinism
    keuzes = np.random.choice(vars['u'].shape[0], aantal_punten, replace=True)
    # calculate the diffusion coordinates # TODO waarom vind Emacs+Python de o-trema niet goed? 
    var_phi1234 = dmap.transform(np.column_stack((vars['gateui'], vars['gatexs'], vars['gated'], vars['Ca'], vars['gatew'], vars['gateu'], vars['Na'], vars['gateoa'], vars['Caup'], vars['gatef'], vars['gateua'], vars['gateh'], vars['gatexr'], vars['gatefCa'], vars['u'], vars['gatem'], vars['gatev'], vars['gatej'], vars['gateoi'], vars['Carel'])))
    print("transformed")
    # plot it!
    fig = plt.figure(figsize=(6,6))
    ax = fig.add_subplot(111,projection='3d')
    ax.scatter(var_phi1234[:,0],var_phi1234[:,1],var_phi1234[:,2],c=var_phi1234[:,3],cmap='viridis')
    ax.set_title("Embedding given by first three DCs, coloured by fourth (Courtemanche model, 1D)") # todo engels
    ax.set_xlabel(r'$\psi_1$')
    ax.set_ylabel(r'$\psi_2$')
    ax.set_zlabel(r'$\psi_3$')
    plt.axis('tight') # ? copied from pydiffmap
    fig.savefig("Courtemanche1D 3+1 eigenvectors, extended embedding.png")
    #fig.savefig("BOFC 3+1 eigenvectors, extended embedding (uvw, s).pdf")
    plt.show()
    plt.close()

matplotlib.use("TkAgg")

# 67: uitstieeksel?
dmap = onderzoek(67, 150, 5200, 120, normalize=True, k=5200)
plt.scatter(range(1,21), dmap.evals)
plt.show()
plt.close()
plot_veel_punten(70,143, 5200, dmap)
